#include "ASMInstruction.hpp"
#include "CodeGen.hpp"
#include "Module.hpp"

#include <iostream>
#include <memory>
#include <unordered_map>

// 函数声明：将高级语言翻译为RISC-V汇编
void translate_main(CodeGen *codegen);
void translate_callee(CodeGen *codegen);

int main() {
    // 初始化编译器模块和代码生成器
    auto *module = new Module();
    auto *codegen = new CodeGen(module);

    // 指定后续生成的汇编代码属于.text段（代码段）
    codegen->append_inst(".text", ASMInstruction::Atrribute);

    // 生成callee和main函数的汇编代码
    translate_callee(codegen);
    translate_main(codegen);

    // 输出生成的汇编代码
    std::cout << codegen->print();
    
    // 清理资源
    delete codegen;
    delete module;
    return 0;
}

// callee函数实现：参数处理和简单运算
void translate_callee(CodeGen *codegen) {
    // 符号表：记录变量在栈帧中的偏移位置
    std::unordered_map<std::string, int> offset_map;

    /* ============== 函数声明部分 ============== */
    codegen->append_inst(".globl callee", ASMInstruction::Atrribute);  // 全局可见
    codegen->append_inst(".type callee, @function", ASMInstruction::Atrribute); // 函数类型

    /* ============== 函数入口点 ============== */
    codegen->append_inst("callee", ASMInstruction::Label);  // 函数入口标签

    /* ============== 序言部分 (Prologue) ============== */
    // 分配48字节栈空间（包含保存寄存器、参数和局部变量的空间）
    codegen->append_inst("addi sp, sp, -48");
    // 保存返回地址寄存器(ra)
    codegen->append_inst("sd ra, 40(sp)");
    // 保存旧的帧指针(s0)
    codegen->append_inst("sd s0, 32(sp)");
    // 设置新的帧指针(s0 = sp + 48)
    codegen->append_inst("addi s0, sp, 48");

    /* ============== 参数处理 ============== */
    offset_map["%arg0"] = -24;  // 第一个参数存储在s0-24位置
    codegen->append_inst("sw a0, -24(s0)");  // 保存参数到栈中

    /* ============== 代码主体部分 ============== */
    codegen->append_inst(".callee_label_entry", ASMInstruction::Label);

    /* 局部变量分配 */
    codegen->append_inst("%op1 = alloca i32", ASMInstruction::Comment);
    // 内存布局：
    // %op1指针变量存储在s0-32
    // 实际分配的i32变量存储在s0-36
    offset_map["%op1"] = -32;
    offset_map["*%op1"] = -36;
    codegen->append_inst("addi t0, s0, -36");  // 计算分配空间的地址
    codegen->append_inst("sd t0, -32(s0)");   // 存储指针到%op1

    /* 参数值存储到局部变量 */
    codegen->append_inst("store i32 %arg0, i32* %op1", ASMInstruction::Comment);
    codegen->append_inst("ld t0, -32(s0)");   // 加载%op1指针
    codegen->append_inst("lw t1, -24(s0)");   // 加载参数值
    codegen->append_inst("sw t1, 0(t0)");     // 存储到%op1指向的位置

    /* 局部变量加载 */
    codegen->append_inst("%op2 = load i32, i32* %op1", ASMInstruction::Comment);
    offset_map["%op2"] = -40;  // 加载结果存储在s0-40
    codegen->append_inst("ld t0, -32(s0)");   // 加载%op1指针
    codegen->append_inst("lw t1, 0(t0)");     // 加载实际值
    codegen->append_inst("sw t1, -40(s0)");    // 存储到%op2

    /* 乘法运算 */
    codegen->append_inst("%op3 = mul i32 3, %op2", ASMInstruction::Comment);
    offset_map["%op3"] = -44;  // 结果存储在s0-44
    codegen->append_inst("lw t0, -40(s0)");    // 加载%op2值
    codegen->append_inst("li t1, 3");          // 加载乘数
    codegen->append_inst("mul t0, t0, t1");    // 执行乘法
    codegen->append_inst("sw t0, -44(s0)");    // 存储结果

    /* ============== 函数返回部分 ============== */
    codegen->append_inst("ret i32 %op3", ASMInstruction::Comment);
    codegen->append_inst("lw a0, -44(s0)");    // 返回值存入a0
    codegen->append_inst("j callee_exit");     // 跳转到公共退出点

    /* ============== 收尾部分 (Epilogue) ============== */
    codegen->append_inst("callee_exit", ASMInstruction::Label);
    codegen->append_inst("ld ra, 40(sp)");     // 恢复返回地址
    codegen->append_inst("ld s0, 32(sp)");     // 恢复帧指针
    codegen->append_inst("addi sp, sp, 48");   // 释放栈空间
    codegen->append_inst("jr ra");             // 函数返回
}

// main函数实现：函数调用和返回值处理
void translate_main(CodeGen *codegen) {
    // 符号表：记录变量在栈帧中的偏移位置
    std::unordered_map<std::string, int> offset_map;

    /* ============== 函数声明部分 ============== */
    codegen->append_inst(".globl main", ASMInstruction::Atrribute);  // 全局可见
    codegen->append_inst(".type main, @function", ASMInstruction::Atrribute); // 函数类型

    /* ============== 函数入口点 ============== */
    codegen->append_inst("main", ASMInstruction::Label);  // 函数入口标签

    /* ============== 序言部分 (Prologue) ============== */
    // 分配32字节栈空间（包含保存寄存器和局部变量的空间）
    codegen->append_inst("addi sp, sp, -32");
    // 保存返回地址寄存器(ra)
    codegen->append_inst("sd ra, 24(sp)");
    // 保存旧的帧指针(s0)
    codegen->append_inst("sd s0, 16(sp)");
    // 设置新的帧指针(s0 = sp + 32)
    codegen->append_inst("addi s0, sp, 32");

    /* ============== 代码主体部分 ============== */
    codegen->append_inst(".main_label_entry", ASMInstruction::Label);

    /* 函数调用 */
    codegen->append_inst("%op0 = call i32 @callee(i32 110)", ASMInstruction::Comment);
    offset_map["%op0"] = -24;  // 返回值存储在s0-24
    codegen->append_inst("li a0, 110");       // 设置参数值
    codegen->append_inst("call callee");      // 调用callee函数
    codegen->append_inst("sw a0, -24(s0)");   // 保存返回值

    /* ============== 函数返回部分 ============== */
    codegen->append_inst("ret i32 %op0", ASMInstruction::Comment);
    codegen->append_inst("lw a0, -24(s0)");   // 加载返回值
    codegen->append_inst("j main_exit");      // 跳转到公共退出点

    /* ============== 收尾部分 (Epilogue) ============== */
    codegen->append_inst("main_exit", ASMInstruction::Label);
    codegen->append_inst("ld ra, 24(sp)");    // 恢复返回地址
    codegen->append_inst("ld s0, 16(sp)");    // 恢复帧指针
    codegen->append_inst("addi sp, sp, 32");  // 释放栈空间
    codegen->append_inst("jr ra");            // 函数返回
}
